#pragma once
#include<iostream>
#include<cassert>
using namespace std;

template<class K, class V>
struct AVLTreeNode {
	//三叉链
	AVLTreeNode<K, V>* _left;
	AVLTreeNode<K, V>* _right;
	AVLTreeNode<K, V>* _parent;

	//平衡因子 Balance factor
	int _bf;

	//存储键值对
	pair<K, V> _kv;
public:
	AVLTreeNode(const pair<K, V>& kv)
		:_left(nullptr)
		,_right(nullptr)
		,_parent(nullptr)
		,_bf(0)
		,_kv(kv)
	{}
};


template<class K, class V>
class AVLTree {
public:
	typedef AVLTreeNode<K, V> Node;
	AVLTree()
		:_root(nullptr)
	{}
	~AVLTree() {
		_Destroy(_root);
	}

public:
	pair<Node*, bool> Insert(const pair<K, V>& kv) {
		//按照二叉搜索树的方式插入新节点

		//_root 为空，直接插入
		if (_root == nullptr) {
			_root = new Node(kv);
			return make_pair(_root, true);
		}
			

		//_root不为空，找到插入位置再插入
		Node* parent = _root;
		Node* cur = _root;

		while (cur) {
			if (cur->_kv.first > kv.first) {
				parent = cur;
				cur = cur->_left;
			}
			else if (cur->_kv.first < kv.first) {
				parent = cur;
				cur = cur->_right;
			}
			else {
				//key已经存在，插入失败
				return make_pair(cur, false);
			}
		}

		//此时cur为空,已经找到插入位置
		
		cur = new Node(kv);
		Node* newNode = cur;

		if (parent->_kv.first > cur->_kv.first) {
			parent->_left = cur;
			cur->_parent = parent;
		}
		else {
			parent->_right = cur;
			cur->_parent = parent;
		}

		//调节节点的平衡因子
		
		while (parent) {
			if (parent->_left == cur) {
				parent->_bf--;
			}
			else {
				parent->_bf++;
			}


			//对AVL树的平衡性进行检查
			if (parent->_bf == 0) {
				break;
			}
			else if (parent->_bf == 1 || parent->_bf == -1) {
				cur = parent;
				parent = parent->_parent;
			}
			else if (parent->_bf == 2 || parent->_bf == -2) {
				//树已经不平衡，旋转

				if (parent->_bf == -2) {
					if (cur->_bf == -1) {
						//右单旋
						_RotateR(parent);
					}
					else { //cur->_bf == 1
						//左右双旋
						_RotateLR(parent);
					}
					
				}
				else { //parent->_bf == 2
					if (cur->_bf == 1) {
						//左单旋
						_RotateL(parent);
					}
					else { //cur->_bf == -1
						//右左双旋
						_RotateRL(parent);
					}
				}
				//旋转完退出
				break;
			}
			else {
				//出错
				assert(false);
			}
		}

		return make_pair(newNode, true);
	}


	Node* Find(const K& key) {
		Node* cur = _root;

		while (cur) {
			if (cur->_kv.first > key) {
				cur = cur->_left;
			}
			else if (cur->_kv.first < key) {
				cur = cur->_right;
			}
			else {
				return cur;
			}
		}

		return nullptr;
	}

	//[]重载
	V& operator[](const K& key) {

		pair<Node*, bool> ret = Insert(make_pair(key, V()));
		return ret.first->_KV.second;
	}

	//中序遍历
	void InOrde() {
		_InOrder(_root);
	}


	bool IsAVLTree() {
		return _IsBalance(_root);
	}


	private:

	//右单旋
	void _RotateR(Node* parent) {
		Node* subL = parent->_left;
		Node* subLR = subL->_right;
		Node* grandParent = parent->_parent;

		parent->_left = subLR;
		//subLR有可能为空
		if (subLR) {
			subLR->_parent = parent;
		}

		subL->_right = parent;
		parent->_parent = subL;


		//看要旋转的点是子树还是根
		if (parent == _root) {
			_root = subL;
			subL->_parent = nullptr;
		}
		else {
			if (grandParent->_left == parent) {
				grandParent->_left = subL;	
			}
			else {
				grandParent->_right = subL;
			}
			subL->_parent = grandParent;
		}


		//更新平衡因子
		subL->_bf = parent->_bf = 0;	
	}

	void _RotateL(Node* parent) {
		Node* subR = parent->_right;
		Node* subRL = subR->_left;
		Node* grandParent = parent->_parent;

		parent->_right = subRL;
		if (subRL) {
			subRL->_parent = parent;
		}

		subR->_left = parent;
		parent->_parent = subR;

		if (parent == _root) {
			_root = subR;
			subR->_parent = nullptr;
		}
		else {
			if (grandParent->_left == parent) {
				grandParent->_left = subR;
			}
			else {
				grandParent->_right = subR;
			}
			subR->_parent = grandParent;
		}

		//更新平衡因子
		subR->_bf = parent->_bf = 0;
	}

	//左右双旋
	void _RotateLR(Node* parent) {
		Node* subL = parent->_left;
		Node* subLR = subL->_right;

		int bf = subLR->_bf;

		_RotateL(subL);
		_RotateR(parent);

		//更新平衡因子
		if (bf == 1) {
			//插入的新结点在subLR右边
			subLR->_bf = 0;
			subL->_bf = -1;
			parent->_bf = 0;
		}
		else if (bf == -1) {
			//插入的新结点在subLR的左边
			subLR->_bf = 0;
			subL->_bf = 0;
			parent->_bf = 1;
		}
		else if (bf == 0) {
			//插入的新结点就是subLR
			subLR->_bf = 0;
			subL->_bf = 0;
			parent->_bf = 0;
		}
		else {
			assert(false);
		}
	}

	//右左双旋
	void _RotateRL(Node* parent) {
		Node* subR = parent->_right;
		Node* subRL = subR->_left;

		int bf = subRL->_bf;

		_RotateR(subR);
		_RotateL(parent);

		if (bf == 1) {
			subRL->_bf = 0; 
			parent->_bf = -1;
			subR->_bf = 0;			
		}
		else if (bf == -1) {
			subRL->_bf = 0;
			parent->_bf = 0;
			subR->_bf = 1;			
		}
		else if (bf == 0) {
			subRL->_bf = 0;
			parent->_bf = 0;
			subR->_bf = 0;			
		}
		else {
			assert(false);
		}
	}

	void _Destroy(Node* root) {
		if (root == nullptr)
			return;

		_Destroy(root->_left);
		_Destroy(root->_right);

		delete root;
	}

	void _InOrder(Node* root) {
		if (root == nullptr) {
			return;
		}

		_InOrder(root->_left);
		cout << root->_kv.first << " : " << root->_kv.second << endl;
		_InOrder(root->_right);

	}

	int _HeightT(Node* root) {
		if (root == nullptr) {
			return 0;
		}

		int leftHeight = _HeightT(root->_left);
		int rightHeight = _HeightT(root->_right);

		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}

	bool _IsBalance(Node* root) {
		if (root == nullptr) {
			return true;
		}

		int leftHeight = _HeightT(root->_left);
		int rightHeight = _HeightT(root->_right);

		if (rightHeight - leftHeight != root->_bf) {
			cout << "平衡因子异常" << root->_kv.first << endl;
			return false;
		}
		
		return abs(rightHeight - leftHeight) < 2
			&& _IsBalance(root->_left)
			&& _IsBalance(root->_right);
	}
	 
private:
	Node* _root;
};